{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ULndVKL12doD"
      },
      "source": [
        "##### Copyright 2019 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "R9SwuuOQE4ko"
      },
      "outputs": [],
      "source": [
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gVgOm1z6DTnV"
      },
      "outputs": [],
      "source": [
        "# These datasets were used in:\n",
        "# \"Are GANs Created Equal? A Large-Scale Study\"\n",
        "# Mario Lucic*, Karol Kurach*, Marcin Michalski, Sylvain Gelly, Olivier Bousquet\n",
        "# NeurIPS 2018\n",
        "\n",
        "# The data set in the paper contains 80000 instances. First 60K are training, next 10K are testing,\n",
        "# and last 10K additional hold-out. You can use this code to create and visualize the datasets.\n",
        "\n",
        "from matplotlib import path\n",
        "import numpy as np\n",
        "from scipy.ndimage import filters\n",
        "\n",
        "def GenerateConvexPolygon(n_vertices, min_segment_angle, scale, raster_dim,\n",
        "                          subpixel_res, shift_to_mean=False):\n",
        "  \"\"\"Returns a random convex polygon with constraints on the minimum angle size.\n",
        "\n",
        "  An ellipse is split into 'n_vertices' segments and a vertex is picked in\n",
        "  each segment. Vertices are then selected within each segment such that the\n",
        "  minimum segment angle (vertex A, zero, vertex B) is respected.\n",
        "\n",
        "  Args:\n",
        "    n_vertices: Integer \u003e 2, number of vertices of the polygon.\n",
        "    min_segment_angle: Float, minimum segment angle (in degrees).\n",
        "    scale: Float in (0, 1], downscaling factor for the generated polygon.\n",
        "    raster_dim: Integer \u003e 1, size of the square grid of pixels to return.\n",
        "    subpixel_res: Integer, each cell of the raster is split into subpixels.\n",
        "    shift_to_mean: Boolean, whether to move the center of mass to the center\n",
        "                   of the raster.\n",
        "  \"\"\"\n",
        "  segment_angle = 360.0 / n_vertices\n",
        "  max_angle = segment_angle - min_segment_angle/2.0\n",
        "  grid_length = raster_dim * scale\n",
        "  angles = []\n",
        "  for i in range(n_vertices):\n",
        "    offset_angle = np.random.rand() * max_angle\n",
        "    angles.append(i * segment_angle + min_segment_angle/2.0 + offset_angle)\n",
        "  angles_in_radians = np.array(angles) * np.pi / 180.0\n",
        "  # Corresponding vertices are mapped to [0, 1]^2.\n",
        "  x = np.column_stack((np.cos(angles_in_radians), np.sin(angles_in_radians)))\n",
        "  x = (x + 1.0) / 2.0\n",
        "\n",
        "  # Randomly rotate the vertices.\n",
        "  theta = np.radians(np.random.rand() * 2 * np.pi)\n",
        "  cos, sin = np.cos(theta), np.sin(theta)\n",
        "  rot = np.matrix([[cos, -sin], [sin, cos]])\n",
        "  x = np.dot(x, rot)\n",
        "\n",
        "  # Corresponding vertices are mapped to [0, raster_dim]^2\n",
        "  vertices = x * grid_length\n",
        "\n",
        "  # If required, move the centre of mass to [raster_dim/2, raster_dim/2]^2.\n",
        "  if shift_to_mean:\n",
        "    vertices += raster_dim / 2.0 - vertices.mean(axis=0)\n",
        "\n",
        "  # Create the bitmask by dividing each cell into subpixels.\n",
        "  r, dim = subpixel_res, raster_dim\n",
        "  poly_as_path = path.Path(vertices * r)\n",
        "  grid_x, grid_y = np.mgrid[0:dim * r, 0:dim * r]\n",
        "  flattened_grid = np.column_stack((grid_x.ravel(), grid_y.ravel()))\n",
        "\n",
        "  mask = poly_as_path.contains_points(flattened_grid).reshape(dim * r, dim * r)\n",
        "  mask = np.array(~mask, dtype=np.float32)\n",
        "  mask = filters.convolve(mask, np.ones((r, r)), mode=\"constant\", cval=1.0)\n",
        "  return mask[::r, ::r] / (r * r)\n",
        "\n",
        "\n",
        "def GenerateDataset(n_instances,\n",
        "                    n_vertices,\n",
        "                    min_segment_angle,\n",
        "                    scale,\n",
        "                    raster_dim,\n",
        "                    subpixel_res,\n",
        "                    shift_to_mean,\n",
        "                    seed=0):\n",
        "  \"\"\"Returns a data set of random convex polygons with angle constraints.\n",
        "\n",
        "  Args:\n",
        "    n_instances: Integer, number of polygons to generate.\n",
        "    n_vertices: Integer \u003e 2, number of vertices of the polygon.\n",
        "    min_segment_angle: Float, minimum angle between two vertices (in degrees).\n",
        "    scale: Float in (0, 1], downscaling factor for the generated polygon.\n",
        "    raster_dim: Integer \u003e 1, size of the square grid of pixels to return.\n",
        "    subpixel_res: Each pixel is split into a grid of size [subpixel_res,\n",
        "                  subpixel_res] and each subpixel is checked when computing\n",
        "                  whether the pixel belongs to the polygon.\n",
        "    shift_to_mean: Boolean, whether to move the center of mass to the center\n",
        "                   of the raster.\n",
        "    seed: Random seed to ensure reproducibility.\n",
        "  Raises:\n",
        "    ValueError: If the number of vertices is too small, the min angle is too\n",
        "                large, or scale is not within (0, 1].\n",
        "  \"\"\"\n",
        "  if n_vertices \u003c 3:\n",
        "    raise ValueError(\"Need more than 2 vertices.\")\n",
        "  if min_segment_angle \u003e 360.0/n_vertices:\n",
        "    raise ValueError(\"The minimum segment angle is infeasible.\")\n",
        "  if scale \u003c 0 or scale \u003e 1:\n",
        "    raise ValueError(\"Scale must be within (0, 1]\")\n",
        "  if raster_dim \u003c= 1:\n",
        "    raise ValueError(\"Raster sidelength has to be greater than 1.\")\n",
        "  np.random.seed(seed)\n",
        "  x = np.zeros((n_instances, raster_dim, raster_dim))\n",
        "  y = np.ones(n_instances, dtype=np.int8) * n_vertices\n",
        "  for i in range(n_instances):\n",
        "    x[i] = GenerateConvexPolygon(\n",
        "        n_vertices=n_vertices,\n",
        "        min_segment_angle=min_segment_angle,\n",
        "        scale=scale,\n",
        "        raster_dim=raster_dim,\n",
        "        subpixel_res=subpixel_res,\n",
        "        shift_to_mean=shift_to_mean)\n",
        "  ids = np.random.permutation(x.shape[0])\n",
        "  return x[ids], y[ids]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "jzj080h6Ddjw"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from mpl_toolkits.axes_grid1 import ImageGrid\n",
        " \n",
        "def visualize(dataset, n_rows):\n",
        "  fig = plt.figure(1, (n_rows, n_rows))\n",
        "  n = dataset.shape[0]\n",
        "  n_cols = n//n_rows + (n % n_rows \u003e 0)\n",
        "  print(n_rows, n_cols)\n",
        "  grid = ImageGrid(fig, 111, nrows_ncols=(n_rows, n_cols), axes_pad=(0, 0))\n",
        "  for i in range(n):\n",
        "    grid[i].set_xlim([0, dataset[0].shape[0]])\n",
        "    grid[i].set_ylim([0, dataset[0].shape[0]])\n",
        "    grid[i].get_xaxis().set_ticks([])\n",
        "    grid[i].get_yaxis().set_ticks([])\n",
        "  for row in range(n_rows):\n",
        "    for col in range(n_cols):\n",
        "      target = dataset[row * n_cols + col]\n",
        "      grid[row * n_cols + col].imshow(target, cmap='gray')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0nH0ozz6D6kQ"
      },
      "outputs": [],
      "source": [
        "triangles, labels = GenerateDataset(\n",
        "    n_instances=100, \n",
        "    n_vertices=3, \n",
        "    min_segment_angle=20, \n",
        "    scale=0.75, \n",
        "    raster_dim=28, \n",
        "    subpixel_res=8, \n",
        "    shift_to_mean=True, \n",
        "    seed=0) \n",
        "visualize(triangles, n_rows=10)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Convex Polygons Dataset",
      "provenance": [
        {
          "file_id": "1NOgTg_rfd_Kp1CK6SA2_PEmrEdOC1Kyc",
          "timestamp": 1568369457874
        }
      ],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
